package com.qm.console.controller.hook.filter.param;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

/**
 * @author fangzhipeng
 * @date 2018/1/8
 */
public class BodyReaderHttpServletResponseWrapper extends HttpServletResponseWrapper {

    private ServletOutputStream servletOutputStream;
    private ByteArrayOutputStream byteArrayOutputStream;
    private PrintWriter printWriter;

    /**
     * Constructs a response adaptor wrapping the given response.
     *
     * @param response The response to be wrapped
     * @throws IllegalArgumentException if the response is null
     */
    public BodyReaderHttpServletResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
        byteArrayOutputStream = new ByteArrayOutputStream();
        servletOutputStream = new MyServletOutputStream(byteArrayOutputStream, response);
        printWriter = new PrintWriter(new OutputStreamWriter(servletOutputStream, "UTF-8"));
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        return printWriter;
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return servletOutputStream;
    }

    public String getResult() throws UnsupportedEncodingException {
        return new String(byteArrayOutputStream.toByteArray(), "UTF-8");
    }

    class MyServletOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream byteArrayOutputStream;
        private HttpServletResponse response;

        public MyServletOutputStream(ByteArrayOutputStream byteArrayOutputStream, HttpServletResponse response) {
            this.byteArrayOutputStream = byteArrayOutputStream;
            this.response = response;
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener listener) {

        }

        @Override
        public void write(int b) throws IOException {
            byteArrayOutputStream.write(b);
            response.getOutputStream().write(b);
        }
    }
}
